import json
import os
from datetime import datetime

from dotenv import load_dotenv
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_community.tools import TavilySearchResults
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import ChatOllama
import matplotlib
from langgraph.prebuilt import create_react_agent
from langgraph.types import interrupt
from pydantic import BaseModel, Field

from langchain_core.tools import tool
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import pymysql
load_dotenv(override=True)
DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
model = ChatOllama(model="qwen3:30b", base_url="http://192.168.97.217:11434")


# 假设 query_result 是数据库返回的字典列表
def default_serializer(obj):
    if isinstance(obj, datetime):
        return obj.strftime('%Y-%m-%d %H:%M:%S')
    raise TypeError(f"Object of type {obj.__class__.__name__} is not JSON serializable")


class WriteQuery(BaseModel):
    content: str = Field(description="需要写入文档的具体内容")
@tool(args_schema=WriteQuery)
def write_file(content):
    """
    将指定的内容写入本地文件
    :param content: 必要参数，字符串类型,用于表示需要写入文档的具体内容
    :return：是否成功写入
    """
    if not content:
        return "内容不能为空"
    file_path = "output.txt"
    with open(file_path, "w") as f:
        f.write(content)
    return "已经成功写入本地文件"


@tool
def human_assistance(query: str) -> str:
    """request assistance from human"""
    human_response = interrupt({"query": query})
    return human_response["data"]


description = """
"当用户需要进行数据库查询工作时，请调用该函数，该函数用于在指定mysql服务器运行一段SQL代码，完成数据查询的相关工作
并且当前函数是使用pymysql连接Mysql数据库.
本函数只负责运行SQL代码并进行数据查询，若要进行数据提取则使用另一个extract_data函数
"""


class SQLQuerySchema(BaseModel):
    sql_query: str = Field(description=description)


@tool(args_schema=SQLQuerySchema)
def sql_inter(sql_query: str) -> str:
    """
    当用户需要进行数据库查询工作时，请调用该函数，该函数用于在指定mysql服务器运行一段SQL代码，完成数据查询的相关工作
    并且当前函数是使用pymysql连接Mysql数据库.
    特别注意mysql日期类型转为python时格式为YYYY-MM-DD HH:MM:SS
    本函数只负责运行SQL代码并进行数据查询，若要进行数据提取则使用另一个extract_data函数
    :param sql_query: 要执行的SQL查询语句
    :return: sql_query在mysql中运行的结果
    """
    load_dotenv(override=True)
    host = os.getenv("MYSQL_HOST")
    user = os.getenv("MYSQL_USER")
    mysql_pw = os.getenv("MYSQL_PW")
    db = os.getenv("MYSQL_DB")
    port = int(os.getenv("MYSQL_PORT"))

    # 连接到MySQL数据库
    connection = pymysql.connect(
        host=host,
        user=user,
        password=mysql_pw,
        database=db,
        port=port,
        charset='utf8'
    )

    try:
        with connection.cursor() as cursor:
            cursor.execute(sql_query)
            result = cursor.fetchall()
            return json.dumps(result, default=default_serializer)

    finally:
        connection.close()


class ExtractQuerySchema(BaseModel):
    sql_query: str = Field(description="用于从Mysql提取数据的SQL查询语句")
    df_name: str = Field(description="用于保存结果的pandas变量名词(字符串形式)")


@tool(args_schema=ExtractQuerySchema)
def extract_data(sql_query: str, df_name: str) -> str:
    """
    用于在Mysql数据库中提取一张表到当前python环境中，注意，本函数只负责数据表的提取，并不负责数据查询。
    若需要在Mysql中进行数据查询，请使用sql_inter函数
    同时需要注意，编写外部函数的参数消息时，必须要满足json格式的字符串""
    :param sql_query: 字符串形式的SQL查询语句，用于提取mysql中的某张表
    :param df_name : 将Mysql数据库中提取的表格进行本地保存时的变量名，以字符串形式表示
    :return: 表格读取和保存结果
    """
    print("正在调用extract_data工具sql查询...")
    load_dotenv(override=True)
    host = os.getenv("MYSQL_HOST")
    user = os.getenv("MYSQL_USER")
    mysql_pw = os.getenv("MYSQL_PW")
    db = os.getenv("MYSQL_DB")
    port = int(os.getenv("MYSQL_PORT"))

    # 连接到MySQL数据库
    connection = pymysql.connect(
        host=host,
        user=user,
        password=mysql_pw,
        database=db,
        port=port,
        charset='utf8'
    )

    try:
        # 执行SQL并保存为全局变量
        df = pd.read_sql(sql_query, connection)
        globals()[df_name] = df
        print(f"数据已保存到变量 {df_name} 中")
        return f"成功创建pandas对象{df_name} ,包含从Mysql提取的数据"
    except Exception as e:
        return f"提取数据时出错: {e}"
    finally:
        connection.close()


class PythonCodeInput(BaseModel):
    py_code: str = Field(description="一段合法的Python 代码字符串，例如'2+1'或 'x=3\\n y=x*2'")


@tool(args_schema=PythonCodeInput)
def python_inter(py_code: str) -> str:
    """
    当用户需要编写python程序并执行时,请调用该函数。
    该函数用于在当前Python环境中执行一段Python代码字符串，并返回执行结果。需要注意本函数只能执行非绘图的代码,
    若是绘图相关代码，则需要调用fig_inter函数运行
    """
    g = globals()
    try:
        return str(eval(py_code, g))
    except Exception as e:
        golobal_vars_before = set(g.keys())
        try:
            exec(py_code, g)
        except Exception as e:
            return f"执行代码报错{e}"
        golobal_vars_after = set(g.keys())
        new_vars = golobal_vars_after - golobal_vars_before
        if new_vars:
            result = {var: g[var] for var in new_vars}
            return str(result)
        else:
            return "代码执行成功"


class FigCodeInput(BaseModel):
    py_code: str = Field(description="要执行的Python绘图代码，必须使用matplotlib/seaborn创建图像并赋值给变量")
    fname: str = Field(description="图像对象的变量名，例如'fig'，用于从代码中提取并保存为图片")


@tool(args_schema=FigCodeInput)
def fig_inter(py_code: str, fname: str) -> str:
    """
    当用户需要编写python进行可视化绘图任务,请调用该函数。
    该函数会执行用户提供的Python绘图代码，并自动将生成的图像对象保存为图片文件并展示
    如果是要运行执行代码而不是绘图，请使用python_inter函数。
    """
    print("正在调用fig_inter工具运行python代码 ")
    current_backend = matplotlib.get_backend()
    matplotlib.use("Agg")

    local_vars = {"plt": plt, "pd": pd, "sns": sns}

    base_dir = '/home/xnzn/disk/workspace/agent-chat-ui/public/images'
    try:
        g = globals()
        exec(py_code, g, local_vars)
        g.update(local_vars)

        fig = local_vars.get(fname, None)
        if fig is None:
            return f"绘图代码中未找到变量名为'{fname}'的图像对象，请检查代码是否正确"
        else:
            fig.savefig(f"{base_dir}/{fname}.png")
            return f"绘图代码执行成功，图像已保存为 {base_dir}/{fname}.png"
    except Exception as e:
        return f"执行绘图代码报错: {e}"


prompt = """
你是一名经验丰富的智能数据分析助手，擅长帮助用户高效完成以下任务：

1. **数据库查询：**
   - 当用户需要获取数据库中的数据或进行 SQL 查询时，请调用 `sql_inter` 工具，该工具已内置 pymysql 连接能力。
   - 你需要准确根据用户请求生成 SQL 语句，例如 `SELECT * FROM 表名` 或包含条件的查询语句。

2. **数据表提取：**
   - 当用户希望将数据库中的表格导入 Python 环境以便后续分析时，请调用 `extract_data` 工具。
   - 根据用户提供的表名或查询条件生成查询语句，并将结果保存到指定的 pandas 变量中。

3. **非绘图类的 Python 代码执行：**
   - 当用户需要执行 Python 脚本进行数据处理、统计计算时，请调用 `python_inter` 工具。
   - 仅执行非绘图类代码，例如变量定义、数据分析、数据清洗等。

4. **绘图类 Python 代码执行：**
   - 当用户需要进行可视化展示（如生成图表、绘制分布图等）时，请调用 `fig_inter` 工具。
   - 可直接读取数据并进行绘图，无需再调用 `python_inter` 工具读取图片。
   - 编写绘图代码时请根据用户需求指定绘图对象变量名（如 `fig`）。
   - 必须使用 `fig = plt.figure()` 或 `fig, ax = plt.subplots()` 显式创建图像对象。
   - 不要调用 `plt.show()`，否则图像将无法保存。

5. **网络搜索：**
   - 当用户提出与数据分析无关的问题（如最新新闻、实时信息），请调用 `search_tool` 工具。

---

**🛠 工具使用优先级：**
- 如需数据库数据，请优先使用 `sql_inter` 或 `extract_data` 获取，再执行 Python 分析或绘图。
- 如需绘图，请确保数据已加载为 pandas 对象。

---

**📌 回答要求：**
- 所有回答均使用 **简体中文**，清晰、礼貌、简洁。
- 如果调用工具返回结构化 JSON 数据，你应提取其中关键字段并进行简要说明，展示主要结论。
- 若需要用户提供更多信息，请主动提出明确的问题。
- 如果生成了图像文件，必须使用 Markdown 格式插入图片，如：
  `![图标题](image/图名.png)`
- 不要仅输出图像路径或文件名。

---

**🎯 回答风格：**
- 专业、简洁，以数据驱动为核心。
- 不要编造不存在的工具或数据。
"""

search = TavilySearchResults(max_results=2)
tools = [search,python_inter,fig_inter,sql_inter,extract_data,write_file,human_assistance]
graph = create_react_agent(model, tools, prompt=prompt)

